from benchmarks.base_benchmark import BaseBenchmark
# -*- coding:utf-8 -*-
import numpy as np
import time

from utils.Ising import Ising


class SBModel():
    """
    A class to implement Ising problems.
    """

    def __to_Ising__(self) -> Ising:
        """
        Generate the equivalent Ising model of the problem.

        Returns
        -------
        Ising
        """

        pass

    def __from_Ising__(self, ising: Ising) -> None:
        """
        Retrieves information from the optimized equivalent Ising model.
        Modifies the object's attributes in place.

        Parameters
        ----------
            ising : Ising
                equivalent Ising model of the problem

        Returns
        -------
        None
        """

        pass


class MaxCut(SBModel):
    """
    A representation of the Markowitz model for portolio optimization.
    Portfolio only takes integer stocks.

    ...

    Attributes
    ----------
    covariance : numpy.ndarray
        the correlation matrix between the assets
    expected_return : numpy.ndarray
        expected return on the investment
    risk_coefficient : float
        risk aversion on the investment
    assets_list : list
        list of the assets to invest in
    number_of_assets : int
        number of different assets
    number_of_bits : int
        number of bits for the binary decomposition of the assets stocks
    portfolio : numpy.ndarray
        array of stocks to purchase per asset
    M : numpy.ndarray
        integer-binary conversion matrix
    ones : numpy.ndarray
        ones vector
    """

    @classmethod
    def from_gset(
            cls,
            gset_path: str
    ):

        """
        Retrieves the data for the Markowitz model from .csv files.
        Only works with the files in the ./data folder.
        """

        with open(gset_path, mode='r') as f:
            title = f.readline()
            if title == 'ROWS':
                print("It is eCut problem.")
                # TODO: process the eCut problem.
                exit(0)
            else:
                vertex_str, edge_str = title.split()
                vertex = int(vertex_str)
                edge = int(edge_str)
                weight = np.zeros((vertex, vertex))
                edge_relationship = f.readlines()
                for cnt in range(edge):
                    edge_a_str, edge_b_str, weight_a_b_str = edge_relationship[cnt].split()
                    edge_a = int(edge_a_str)
                    edge_b = int(edge_b_str)
                    weight_a_b = int(weight_a_b_str)
                    weight[edge_a - 1][edge_b - 1] = weight_a_b
                    weight[edge_b - 1][edge_a - 1] = weight_a_b

        return MaxCut(
            vertex,
            edge,
            weight
        )

    def __init__(
            self,
            vertex: int,
            edge: int,
            weight: np.ndarray,
            assert_parameters: int = 1
    ) -> None:

        """
        Constructs all the necessary attributes for the Markowitz object.

        Parameters
        ----------
            vertex : int
                the vertex num of
            expected_return : numpy.ndarray
                expected return on the investment
            risk_coefficient : float
                risk aversion on the investment
            assets_list : list
                list of the assets to invest in
            number_of_bits : int
                number of bits for the binary decomposition of the assets stocks
            assert_parameters : bool, optional
                check the format of the inputs (default is True)
        """

        # Data
        self.vertex = vertex
        self.edge = edge
        self.weight = weight
        self.cut_value = vertex
        if assert_parameters:
            self.__assert__()

    def __repr__(self) -> str:

        return f"""Utility gain: {self.utlity_function()}
        {self.as_dataframe()}
        """

    def __assert__(self) -> None:

        """
        Checks the format of the attributes.

        Returns
        -------
        float
        """

        # Checking types

        # Checking dimensions

        pass

    def __to_Ising__(self) -> Ising:

        J = -self.weight
        h = np.zeros((self.vertex, 1))

        return Ising(J, h)

    def __from_Ising__(self, ising: Ising) -> None:
        if self.cut_value <= -0.25 * np.sum(
                -self.weight) - 0.25 * ising.ground_state.T @ self.weight @ ising.ground_state:
            self.lowest_ground_state = ising.ground_state.copy()
            self.cut_value = -0.25 * np.sum(
                -self.weight) - 0.25 * ising.ground_state.T @ self.weight @ ising.ground_state


class Benchmark(BaseBenchmark):

    def __init__(self, config):
        super().__init__(config)
        gset_path = "./data/Gset/G22.txt"
        self.maxcut = MaxCut.from_gset(gset_path)
        self.ising_equivalent = self.maxcut.__to_Ising__()

        assert config['n_dim'] == self.ising_equivalent.dimension

    def evaluate(self, x):
        self.ising_equivalent.get_ground_state(x[:, None])
        return self.ising_equivalent.energy()
